# encoding=utf8
import os
import json
import re
import math
import codecs
import requests
from pymongo import MongoClient
from bson import json_util
from urllib.parse import quote_plus
from tqdm import tqdm


image_names = ['单行公式', '多行公式', '剩余公式', '不确定公式']
image_paths = {name: os.path.join('./gongshi_image/{}'.format(name)) for name in image_names}
for img_pth in image_paths.values():
    if not os.path.exists(img_pth):
        os.makedirs(img_pth)

danhang_count, duohang_count, shengyu_count, buqueding_count = 0, 0, 0, 0

user = 'mongouser'
pwd = 'tuanjie123'
host = '118.25.106.14'
port = 27017
url = 'mongodb://%s:%s@%s' % (quote_plus(user), quote_plus(pwd), host)

client = MongoClient(url, port=port)
collection_task_test = client['tagging_test']['task_test']

# db.getCollection('task_test').find({taskKey: 'sentimentAnalysis', canOutput: 1, checkDate: {'$gte': 1654704000000, '$lte': 1654790400000}}).count()
lines = collection_task_test.find(
    {
        'taskKey': 'formulaCate',
        'canOutput': 1,
        # 'checkDate': {'$gte': 1658332800000}  # 8-4: 1659542400000
    },
    {
        'taskData': 1,
        'checkData': 1,
        # 'reCheckData': 1,
        # 'finalData': 1,
        '_id': 0,
    },
    no_cursor_timeout=True
)

total = lines.count()
print(total)

success_count = 0
error_count = 0

for idx, line in tqdm(enumerate(lines), total=total):  # lines是动态读取mongo中的数据
    file_name = line['checkData']['fileName']
    image_tag = line['checkData']['formulaType']
    if (image_tag == '单行公式' and danhang_count >= 3000) or \
        (image_tag == '多行公式' and duohang_count >= 3000) or \
        (image_tag == '不确定公式' and buqueding_count >= 3000) \
        or (image_tag == '剩余公式' and shengyu_count >= 3000):
        continue
    try:
        req_url = f'https://datafactory.xind.com/fs/static/{file_name}'
        res = requests.get(req_url)
        # 存储图片
        with open(os.path.join(image_paths[image_tag], file_name), 'wb') as img_f:
            img_f.write(res.content)

        if image_tag == '单行公式':
            danhang_count += 1
        if image_tag == '多行公式':
            duohang_count += 1
        if image_tag == '不确定公式':
            buqueding_count += 1
        if image_tag == '剩余公式':
            shengyu_count += 1
        # # 存储标签
        # with open(os.path.join(image_paths[image_tag], file_name.split('.')[0] + '.txt'), 'w', encoding='utf-8') as txt_f:
        #     txt_f.write(image_tag)

        success_count += 1

        # if idx == 50:
        #     break

    except:
        error_count += 1

print(danhang_count, duohang_count, shengyu_count, buqueding_count)
print(success_count)
print(error_count)

